{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quantized RNNs and LSTMs\n",
    "\n",
    "With version 0.8, Brevitas introduces support for quantized recurrent layers through `QuantRNN` and `QuantLSTM`. As with other Brevitas quantized layers, `QuantRNN` and `QuantLSTM` can be used as drop-in replacement for their floating-point variants, but they also go further and support some additional structural recurrent options not found in upstream PyTorch. Similarly to other quantized layers, both `QuantRNN` and `QuantLSTM` can take in different quantizers for different tensors involved in their computation."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## QuantRNN\n",
    "\n",
    "We start by looking at `QuantRNN`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:43.661853Z",
     "iopub.status.busy": "2025-05-09T14:50:43.661158Z",
     "iopub.status.idle": "2025-05-09T14:50:50.386785Z",
     "shell.execute_reply": "2025-05-09T14:50:50.384321Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```python\n",
       "    def __init__(\n",
       "            self,\n",
       "            input_size: int,\n",
       "            hidden_size: int,\n",
       "            num_layers: int = 1,\n",
       "            nonlinearity: str = 'tanh',\n",
       "            bias: Optional[bool] = True,\n",
       "            batch_first: bool = False,\n",
       "            bidirectional: bool = False,\n",
       "            weight_quant=Int8WeightPerTensorFloat,\n",
       "            bias_quant=Int32Bias,\n",
       "            io_quant=Int8ActPerTensorFloat,\n",
       "            gate_acc_quant=Int8ActPerTensorFloat,\n",
       "            shared_input_hidden_weights=False,\n",
       "            return_quant_tensor: bool = False,\n",
       "            dtype: Optional[torch.dtype] = None,\n",
       "            device: Optional[torch.device] = None,\n",
       "            **kwargs):\n",
       "        super(QuantRNN, self).__init__(\n",
       "            layer_impl=_QuantRNNLayer,\n",
       "            input_size=input_size,\n",
       "            hidden_size=hidden_size,\n",
       "            num_layers=num_layers,\n",
       "            nonlinearity=nonlinearity,\n",
       "            bias=bias,\n",
       "            batch_first=batch_first,\n",
       "            bidirectional=bidirectional,\n",
       "            weight_quant=weight_quant,\n",
       "            bias_quant=bias_quant,\n",
       "            io_quant=io_quant,\n",
       "            gate_acc_quant=gate_acc_quant,\n",
       "            shared_input_hidden_weights=shared_input_hidden_weights,\n",
       "            return_quant_tensor=return_quant_tensor,\n",
       "            dtype=dtype,\n",
       "            device=device,\n",
       "            **kwargs)\n",
       "\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import inspect\n",
    "from brevitas.nn import QuantRNN\n",
    "from IPython.display import Markdown, display\n",
    "import torch\n",
    "torch.manual_seed(0)\n",
    "\n",
    "# helpers\n",
    "def assert_with_message(condition):\n",
    "    assert condition\n",
    "    print(condition)\n",
    "\n",
    "def pretty_print_source(source):\n",
    "    display(Markdown('```python\\n' + source + '\\n```'))\n",
    "    \n",
    "source = inspect.getsource(QuantRNN.__init__)  \n",
    "pretty_print_source(source)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`QuantRNN` supports all arguments of `torch.nn.RNN`, plus it exposes four different quantizers: `weight_quant` controls quantization of the weight tensor, `bias_quant` controls quantization of the bias, `io_quant` controls quantization of the input/output, and `gate_acc_quant` controls quantization of the output of the gate, before the nonlinearity is applied. \n",
    "\n",
    "Compared to other layers like `QuantLinear`, a couple of things can be observed. First, input and output quantization are fused together into `io_quant`. This is because of the recurrent structure of RNN layers, where the output is fed back as input. Second, all quantizers are already set by default. This is different from a layer like `QuantLinear`, where only `weight_quant` has a default quantizer.\n",
    "\n",
    "As with `torch.nn.RNN`, `QuantRNN` defines a stack of potentially multiple layers, controlled by setting `num_layers`, that can be set to bidirectional with `bidirectional=True`. Internally, `QuantRNN` is organized into a two level nesting of `ModuleList`, one for the different layer(s), and one for the direction(s):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:50.547938Z",
     "iopub.status.busy": "2025-05-09T14:50:50.547053Z",
     "iopub.status.idle": "2025-05-09T14:50:50.736072Z",
     "shell.execute_reply": "2025-05-09T14:50:50.733601Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/nn/mixin/base.py:59: UserWarning: Keyword arguments are being passed but they not being used.\n",
      "  warn('Keyword arguments are being passed but they not being used.')\n"
     ]
    }
   ],
   "source": [
    "def rnn_sublayer(module, sublayer_number, right_to_left_direction):\n",
    "    return module.layers[sublayer_number][1 if right_to_left_direction else 0]\n",
    "\n",
    "quant_rnn = QuantRNN(input_size=10, hidden_size=20, num_layers=2, bidirectional=True)\n",
    "quant_rnn_0_left_to_right = rnn_sublayer(quant_rnn, sublayer_number=0, right_to_left_direction=False)\n",
    "quant_rnn_0_right_to_left = rnn_sublayer(quant_rnn, sublayer_number=0, right_to_left_direction=True)\n",
    "quant_rnn_1_left_to_right = rnn_sublayer(quant_rnn, sublayer_number=1, right_to_left_direction=False)\n",
    "quant_rnn_1_right_to_left = rnn_sublayer(quant_rnn, sublayer_number=1, right_to_left_direction=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setting `num_layers > 1` and/or `bidirectional=True` has different implications on different quantizers. For `weight_quant`, `gate_acc_quant` and `bias_quant`, the same quantizer *definition* is shared among different layers/directions, but each layer/direction is allocated its own instance of the quantizer. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:50.746819Z",
     "iopub.status.busy": "2025-05-09T14:50:50.745414Z",
     "iopub.status.idle": "2025-05-09T14:50:50.757818Z",
     "shell.execute_reply": "2025-05-09T14:50:50.755397Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "assert_with_message(not quant_rnn_0_left_to_right.gate_params.input_weight.weight_quant is quant_rnn_1_right_to_left.gate_params.input_weight.weight_quant)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:50.766388Z",
     "iopub.status.busy": "2025-05-09T14:50:50.765712Z",
     "iopub.status.idle": "2025-05-09T14:50:50.777230Z",
     "shell.execute_reply": "2025-05-09T14:50:50.775002Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "assert_with_message(not quant_rnn_0_left_to_right.cell.gate_acc_quant is quant_rnn_1_right_to_left.cell.gate_acc_quant)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:50.785905Z",
     "iopub.status.busy": "2025-05-09T14:50:50.785020Z",
     "iopub.status.idle": "2025-05-09T14:50:50.796842Z",
     "shell.execute_reply": "2025-05-09T14:50:50.794454Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "assert_with_message(not quant_rnn_0_left_to_right.gate_params.bias_quant is quant_rnn_1_right_to_left.gate_params.bias_quant)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Conversely, for `io_quant` the same *instance* is gonna be shared among all layers and directions. This is to make sure that input/output tensors that are internally concatenated together share the same quantization scale/zero-point/bitwidth. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:50.806419Z",
     "iopub.status.busy": "2025-05-09T14:50:50.805604Z",
     "iopub.status.idle": "2025-05-09T14:50:50.817169Z",
     "shell.execute_reply": "2025-05-09T14:50:50.814930Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "assert_with_message(quant_rnn_0_left_to_right.io_quant is quant_rnn_1_right_to_left.io_quant)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, `QuantRNN` supports an additional flag, `shared_input_hidden_weights`. This allows, whenever `bidirectional=True`, to share the input-to-hidden weights among the two directions, an optimization introduced first by DeepSpeech back in the day to save on the number of parameters, with minimal impact on the quality of results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:50.826929Z",
     "iopub.status.busy": "2025-05-09T14:50:50.825968Z",
     "iopub.status.idle": "2025-05-09T14:50:51.045372Z",
     "shell.execute_reply": "2025-05-09T14:50:51.042919Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of weights for single direction QuantRNN: 600\n",
      "Number of weights for bidirectional QuantRNN: 1200\n",
      "Number of weights for bidirectional QuantRNN with shared input-hidden weights: 1000\n"
     ]
    }
   ],
   "source": [
    "from brevitas.nn import QuantRNN\n",
    "\n",
    "def count_weights(model):\n",
    "    return sum(p.numel() for n, p in model.named_parameters() if 'weight' in n)\n",
    "\n",
    "quant_rnn_single_direction = QuantRNN(input_size=10, hidden_size=20, bidirectional=False, shared_input_hidden_weights=False)\n",
    "quant_rnn_bidirectional = QuantRNN(input_size=10, hidden_size=20, bidirectional=True, shared_input_hidden_weights=False)\n",
    "quant_rnn_bidirectional_shared_input_hidden = QuantRNN(input_size=10, hidden_size=20, bidirectional=True, shared_input_hidden_weights=True)\n",
    "\n",
    "print(f\"Number of weights for single direction QuantRNN: {count_weights(quant_rnn_single_direction)}\")\n",
    "print(f\"Number of weights for bidirectional QuantRNN: {count_weights(quant_rnn_bidirectional)}\")\n",
    "print(f\"Number of weights for bidirectional QuantRNN with shared input-hidden weights: {count_weights(quant_rnn_bidirectional_shared_input_hidden)}\")\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As with other Brevitas layers, it's possible to directly modify a quantizer by passing keyword arguments with a matching prefix. For example, to set 4b per-channel weights and 6b io quantization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:51.054690Z",
     "iopub.status.busy": "2025-05-09T14:50:51.053973Z",
     "iopub.status.idle": "2025-05-09T14:50:51.142152Z",
     "shell.execute_reply": "2025-05-09T14:50:51.139886Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input-hidden weight bit-width: 4.0\n",
      "Hidden-hidden weight bit-width: 4.0\n",
      "I/O quant bit-width: 6.0\n",
      "Input-hidden weight scale: tensor([[0.0297],\n",
      "        [0.0311],\n",
      "        [0.0298],\n",
      "        [0.0295],\n",
      "        [0.0316],\n",
      "        [0.0311],\n",
      "        [0.0318],\n",
      "        [0.0309],\n",
      "        [0.0317],\n",
      "        [0.0309],\n",
      "        [0.0316],\n",
      "        [0.0319],\n",
      "        [0.0319],\n",
      "        [0.0318],\n",
      "        [0.0315],\n",
      "        [0.0310],\n",
      "        [0.0319],\n",
      "        [0.0319],\n",
      "        [0.0318],\n",
      "        [0.0312]], grad_fn=<DivBackward0>)\n",
      "Hidden-hidden weight scale: tensor([[0.0297],\n",
      "        [0.0311],\n",
      "        [0.0298],\n",
      "        [0.0295],\n",
      "        [0.0316],\n",
      "        [0.0311],\n",
      "        [0.0318],\n",
      "        [0.0309],\n",
      "        [0.0317],\n",
      "        [0.0309],\n",
      "        [0.0316],\n",
      "        [0.0319],\n",
      "        [0.0319],\n",
      "        [0.0318],\n",
      "        [0.0315],\n",
      "        [0.0310],\n",
      "        [0.0319],\n",
      "        [0.0319],\n",
      "        [0.0318],\n",
      "        [0.0312]], grad_fn=<DivBackward0>)\n"
     ]
    }
   ],
   "source": [
    "quant_rnn_4b = QuantRNN(input_size=10, hidden_size=20, weight_bit_width=4, weight_scaling_per_output_channel=True, io_bit_width=6)\n",
    "quant_rnn_4b_0_left_to_right = rnn_sublayer(quant_rnn_4b, sublayer_number=0, right_to_left_direction=False)\n",
    "\n",
    "input_hidden_weight = quant_rnn_4b_0_left_to_right.gate_params.input_weight.quant_weight()\n",
    "hidden_hidden_weight = quant_rnn_4b_0_left_to_right.gate_params.hidden_weight.quant_weight()\n",
    "\n",
    "print(f\"Input-hidden weight bit-width: {input_hidden_weight.bit_width}\")\n",
    "print(f\"Hidden-hidden weight bit-width: {hidden_hidden_weight.bit_width}\")\n",
    "print(f\"I/O quant bit-width: {quant_rnn_4b_0_left_to_right.io_quant.bit_width()}\")\n",
    "print(f\"Input-hidden weight scale: {input_hidden_weight.scale}\")\n",
    "print(f\"Hidden-hidden weight scale: {hidden_hidden_weight.scale}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`QuantRNN` follows the same `forward` interface of `torch.nn.RNN`, with a couple of exceptions. Packed variable length inputs are currently not supported, and unbatched inputs are not supported. \n",
    "Other than that, everything else is the same. \n",
    "\n",
    "Inputs are expected to have shape `(batch, sequence, features)` for `batch_first=False`, or `(sequence, batch, features)` for `batch_first=True`. The layer returns a tuple with `(outputs, hidden_states)`, where `outputs` has shape `(sequence, batch, hidden_size * num_directions)` with `num_directions=2` when `bidirectional=True`, for `batch_first=False`, or `(batch, sequence, hidden_size * num_directions)` for `batch_first=True`, while `hidden_states` has shape `(num_directions * num_layers, batch, hidden_size)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:51.153953Z",
     "iopub.status.busy": "2025-05-09T14:50:51.152734Z",
     "iopub.status.idle": "2025-05-09T14:50:51.248126Z",
     "shell.execute_reply": "2025-05-09T14:50:51.245439Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output size: torch.Size([2, 5, 20])\n",
      "Hidden states size: torch.Size([1, 2, 20])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from brevitas.nn import QuantRNN\n",
    "\n",
    "quant_rnn = QuantRNN(input_size=10, hidden_size=20, batch_first=True)\n",
    "outputs, hidden_states = quant_rnn(torch.randn(2, 5, 10))\n",
    "print(f\"Output size: {outputs.shape}\")\n",
    "print(f\"Hidden states size: {hidden_states.shape}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As with other quantized layers, it's possible to return a `QuantTensor` with `return_quant_tensor=True`. As a reminder, a `QuantTensor` is just a data structure that captures the quantization metadata associated with a quantized tensor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:51.259788Z",
     "iopub.status.busy": "2025-05-09T14:50:51.258937Z",
     "iopub.status.idle": "2025-05-09T14:50:51.369690Z",
     "shell.execute_reply": "2025-05-09T14:50:51.367400Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(IntQuantTensor(value=tensor([[[-0.0062, -0.2872,  0.7931,  0.4309,  0.5495, -0.4558,  0.2373,\n",
       "            0.6807,  0.4621,  0.6120, -0.1124,  0.3872,  0.3060,  0.7681,\n",
       "           -0.3684,  0.0437, -0.7369, -0.3247,  0.7743,  0.3372],\n",
       "          [ 0.5450,  0.2962, -0.3969,  0.3555, -0.5628,  0.2429, -0.4976,\n",
       "            0.1777, -0.1244,  0.0296, -0.2607,  0.0948,  0.5036, -0.3673,\n",
       "            0.5213, -0.2962,  0.7524,  0.0770, -0.0948, -0.0948],\n",
       "          [ 0.2691, -0.6624, -0.5434,  0.4968, -0.6624,  0.0983,  0.1345,\n",
       "            0.1242, -0.0517, -0.3726,  0.3053,  0.1604,  0.3208,  0.0983,\n",
       "            0.3105,  0.4243,  0.2794,  0.1604,  0.1035, -0.0724],\n",
       "          [ 0.1284, -0.3337, -0.5263, -0.0449, -0.5263,  0.3081, -0.1733,\n",
       "            0.5648,  0.4942, -0.1412,  0.1733,  0.3337,  0.6225,  0.3401,\n",
       "            0.5070, -0.1412,  0.0642, -0.3722,  0.2888,  0.1155],\n",
       "          [ 0.0579, -0.0058, -0.4054, -0.1564, -0.5560, -0.3301,  0.3533,\n",
       "            0.0058, -0.1622, -0.3765,  0.1216,  0.0695, -0.4054,  0.0927,\n",
       "            0.6139, -0.1390,  0.7066,  0.1274,  0.1622, -0.2896]],\n",
       " \n",
       "         [[ 0.1374,  0.5745,  0.0624, -0.2373,  0.3060,  0.3310, -0.5183,\n",
       "            0.1186,  0.1124,  0.2997,  0.0375,  0.6369, -0.5308,  0.6307,\n",
       "           -0.5683,  0.7556,  0.2997, -0.4933,  0.3934, -0.4871],\n",
       "          [ 0.1066, -0.1244, -0.1718,  0.4266,  0.5569,  0.0178,  0.1185,\n",
       "           -0.3910,  0.2133,  0.0178, -0.1066, -0.2903,  0.1837, -0.2547,\n",
       "           -0.2903,  0.0770,  0.3495,  0.2547,  0.2311, -0.6161],\n",
       "          [-0.0880, -0.1966,  0.3001, -0.0569,  0.4140, -0.1552, -0.1345,\n",
       "            0.4554,  0.5175,  0.1242, -0.2898,  0.1966, -0.0414,  0.3985,\n",
       "           -0.1708, -0.0621, -0.1708,  0.0828,  0.2225,  0.0517],\n",
       "          [ 0.2118,  0.5648, -0.2824, -0.0449,  0.5840,  0.3209, -0.5648,\n",
       "            0.3530,  0.4043, -0.4942, -0.3786,  0.0257,  0.5327, -0.1990,\n",
       "           -0.1348, -0.8215,  0.3016,  0.5327,  0.5648, -0.1155],\n",
       "          [-0.0290, -0.1738,  0.0695,  0.3765,  0.1738,  0.0579, -0.4054,\n",
       "           -0.2664,  0.4923,  0.2143, -0.4170,  0.4112,  0.5502,  0.7066,\n",
       "           -0.6024,  0.7356,  0.0348,  0.1043, -0.1911, -0.4518]]],\n",
       "        grad_fn=<CatBackward0>), scale=tensor(0.0059, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True)),\n",
       " IntQuantTensor(value=tensor([[[ 0.0579, -0.0058, -0.4054, -0.1564, -0.5560, -0.3301,  0.3533,\n",
       "            0.0058, -0.1622, -0.3765,  0.1216,  0.0695, -0.4054,  0.0927,\n",
       "            0.6139, -0.1390,  0.7066,  0.1274,  0.1622, -0.2896],\n",
       "          [-0.0290, -0.1738,  0.0695,  0.3765,  0.1738,  0.0579, -0.4054,\n",
       "           -0.2664,  0.4923,  0.2143, -0.4170,  0.4112,  0.5502,  0.7066,\n",
       "           -0.6024,  0.7356,  0.0348,  0.1043, -0.1911, -0.4518]]],\n",
       "        grad_fn=<UnsqueezeBackward0>), scale=tensor(0.0058, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True)))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from brevitas.nn import QuantRNN\n",
    "\n",
    "quant_rnn = QuantRNN(input_size=10, hidden_size=20, batch_first=True, return_quant_tensor=True)\n",
    "quant_rnn(torch.randn(2, 5, 10))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly, a `QuantTensor` can be passed in as input. However, whenever `io_quant` is set (which it is by default), the input is gonna be re-quantized:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:51.378775Z",
     "iopub.status.busy": "2025-05-09T14:50:51.378070Z",
     "iopub.status.idle": "2025-05-09T14:50:51.460239Z",
     "shell.execute_reply": "2025-05-09T14:50:51.458021Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/torch/_tensor.py:1255: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at /opt/conda/conda-bld/pytorch_1670525541990/work/c10/core/TensorImpl.h:1758.)\n",
      "  return super(Tensor, self).rename(names)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(IntQuantTensor(value=tensor([[[ 0.2111,  0.1267,  0.0060,  0.6153, -0.7721, -0.3740, -0.5188,\n",
       "            0.6273,  0.4162,  0.2051,  0.2292,  0.7239,  0.6032,  0.2533,\n",
       "            0.5067,  0.6635,  0.1206, -0.5730,  0.0483,  0.3318],\n",
       "          [ 0.5742,  0.0194, -0.3807, -0.0710, -0.6000,  0.1807,  0.1355,\n",
       "            0.4129,  0.3807,  0.3936, -0.0903,  0.1549,  0.1032,  0.0645,\n",
       "            0.4775, -0.0645,  0.1161, -0.0065,  0.0194, -0.1097],\n",
       "          [ 0.0453, -0.4533,  0.1036, -0.0194, -0.2979,  0.3432,  0.0777,\n",
       "            0.6346, -0.0842,  0.3302,  0.4727,  0.4856, -0.4144,  0.7382,\n",
       "           -0.0453,  0.5439,  0.2266, -0.4792,  0.4403, -0.1036],\n",
       "          [ 0.3198,  0.2741, -0.6395,  0.0971, -0.6052, -0.5196,  0.1770,\n",
       "           -0.5025, -0.1256,  0.2056,  0.2684, -0.6395, -0.0285, -0.7309,\n",
       "            0.7194, -0.7194,  0.1542, -0.3426, -0.6509,  0.0343],\n",
       "          [ 0.0000, -0.4004,  0.3151, -0.0263, -0.5842, -0.1641, -0.3939,\n",
       "            0.0263, -0.2429,  0.6499, -0.5186,  0.1247, -0.2101,  0.8337,\n",
       "           -0.1444,  0.6762, -0.1641, -0.5317, -0.1707, -0.0197]],\n",
       " \n",
       "         [[ 0.2111, -0.2111, -0.3197, -0.0241, -0.5067, -0.0241, -0.2895,\n",
       "            0.1749, -0.4283,  0.0000, -0.3680,  0.5308, -0.1267,  0.5248,\n",
       "            0.1206,  0.2654,  0.6394, -0.1327,  0.2292, -0.3800],\n",
       "          [ 0.6775, -0.3355, -0.1807,  0.2774, -0.8259, -0.2000, -0.0065,\n",
       "            0.5678,  0.4000,  0.2258,  0.4387,  0.2710,  0.5355,  0.1290,\n",
       "            0.6710, -0.0645, -0.2710, -0.3613,  0.6388,  0.5226],\n",
       "          [-0.0065, -0.0777, -0.6475, -0.1684, -0.3820,  0.3885,  0.0065,\n",
       "            0.1943, -0.3238, -0.2525, -0.1230, -0.0453, -0.0777,  0.3432,\n",
       "            0.4921, -0.1101,  0.8224,  0.2396,  0.1554, -0.3885],\n",
       "          [-0.0514, -0.4111, -0.4625, -0.1713, -0.3369,  0.2512, -0.2969,\n",
       "           -0.4111, -0.2341,  0.3597, -0.1998,  0.0000,  0.2741,  0.7137,\n",
       "           -0.1256,  0.1370, -0.0742, -0.5938, -0.5424, -0.4168],\n",
       "          [ 0.3479,  0.5974, -0.3939,  0.1444, -0.6762,  0.1969, -0.6499,\n",
       "            0.4136,  0.5383, -0.3085,  0.4070,  0.4070,  0.6630, -0.0263,\n",
       "            0.2823, -0.1510,  0.1313, -0.5186,  0.4464, -0.0066]]],\n",
       "        grad_fn=<CatBackward0>), scale=tensor(0.0062, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True)),\n",
       " IntQuantTensor(value=tensor([[[ 0.0000, -0.4004,  0.3151, -0.0263, -0.5842, -0.1641, -0.3939,\n",
       "            0.0263, -0.2429,  0.6499, -0.5186,  0.1247, -0.2101,  0.8337,\n",
       "           -0.1444,  0.6762, -0.1641, -0.5317, -0.1707, -0.0197],\n",
       "          [ 0.3479,  0.5974, -0.3939,  0.1444, -0.6762,  0.1969, -0.6499,\n",
       "            0.4136,  0.5383, -0.3085,  0.4070,  0.4070,  0.6630, -0.0263,\n",
       "            0.2823, -0.1510,  0.1313, -0.5186,  0.4464, -0.0066]]],\n",
       "        grad_fn=<UnsqueezeBackward0>), scale=tensor(0.0066, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True)))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from brevitas.nn import QuantIdentity\n",
    "\n",
    "quant_identity = QuantIdentity(return_quant_tensor=True)\n",
    "quant_rnn(quant_identity(torch.randn(2, 5, 10)))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As with `torch.nn.RNN`, by default the initial hidden state is initialized to 0, but a custom hidden state of shape `(num_directions * num_layers, batch, hidden_size)` can be passed in:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:51.469570Z",
     "iopub.status.busy": "2025-05-09T14:50:51.468766Z",
     "iopub.status.idle": "2025-05-09T14:50:51.535635Z",
     "shell.execute_reply": "2025-05-09T14:50:51.533370Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(IntQuantTensor(value=tensor([[[-0.3777, -0.2074,  0.7184,  0.9110,  0.0148, -0.1926, -0.7110,\n",
       "            0.1926, -0.4222, -0.9480,  0.2592,  0.2222, -0.2370, -0.5407,\n",
       "            0.5851, -0.2370,  0.3555,  0.1703,  0.4444, -0.2222],\n",
       "          [ 0.4814, -0.7355, -0.1605,  0.3878, -0.5282,  0.2073,  0.0000,\n",
       "            0.3677,  0.1805, -0.1204, -0.4614,  0.2474,  0.7021,  0.0401,\n",
       "            0.4346,  0.4480, -0.3143,  0.0401,  0.6887,  0.6753],\n",
       "          [ 0.5038, -0.3650, -0.6936,  0.0146, -0.9345,  0.0000,  0.1679,\n",
       "           -0.3066,  0.1825,  0.4089,  0.0949, -0.2555,  0.3870, -0.2482,\n",
       "            0.5914, -0.0803,  0.1314, -0.4235, -0.3797,  0.1168],\n",
       "          [ 0.1795,  0.1795,  0.0449,  0.0449,  0.2308,  0.0898, -0.1282,\n",
       "            0.5579,  0.1731, -0.1795,  0.1603,  0.3142,  0.1090,  0.5835,\n",
       "           -0.1475,  0.0449,  0.1795, -0.0256,  0.8143, -0.2437],\n",
       "          [-0.0066,  0.4804,  0.0066, -0.1184,  0.6843, -0.0197,  0.1448,\n",
       "            0.1842,  0.6383, -0.1908, -0.0066, -0.1053, -0.1316,  0.0461,\n",
       "           -0.0066, -0.2764,  0.3751,  0.3619,  0.5001, -0.1316]],\n",
       " \n",
       "         [[ 0.5110, -0.3555,  0.6443, -0.8221,  0.4888, -0.2074,  0.0444,\n",
       "            0.4888,  0.5999,  0.4370,  0.0000,  0.5036, -0.7628,  0.9332,\n",
       "           -0.6147,  0.7332,  0.3629,  0.9184,  0.7702, -0.8887],\n",
       "          [ 0.8492, -0.3410, -0.3878,  0.1404, -0.3410,  0.3143, -0.1204,\n",
       "            0.5817,  0.4413,  0.5550,  0.6486, -0.1070,  0.6285, -0.4948,\n",
       "            0.2006,  0.1605,  0.0535, -0.4079,  0.3811,  0.4948],\n",
       "          [ 0.6060,  0.7666, -0.8688, -0.6863, -0.5111, -0.0803, -0.6425,\n",
       "           -0.0146, -0.3577,  0.3431, -0.6571,  0.5622,  0.0000,  0.7374,\n",
       "           -0.1314, -0.3650,  0.7520,  0.2336, -0.2847, -0.8250],\n",
       "          [ 0.3014,  0.2950, -0.0898, -0.3142,  0.4040,  0.4681, -0.0705,\n",
       "           -0.2052,  0.8143, -0.1603,  0.3334, -0.6733,  0.0834,  0.0898,\n",
       "           -0.4937,  0.1924,  0.0064,  0.4104,  0.6348, -0.3527],\n",
       "          [-0.6449,  0.5856, -0.0263, -0.0197,  0.8357, -0.5856,  0.0395,\n",
       "           -0.3422,  0.8028,  0.0855, -0.7238, -0.6317,  0.2764, -0.0461,\n",
       "           -0.4211, -0.5988,  0.2632,  0.4014, -0.7501, -0.5659]]],\n",
       "        grad_fn=<CatBackward0>), scale=tensor(0.0069, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True)),\n",
       " IntQuantTensor(value=tensor([[[-0.0066,  0.4804,  0.0066, -0.1184,  0.6843, -0.0197,  0.1448,\n",
       "            0.1842,  0.6383, -0.1908, -0.0066, -0.1053, -0.1316,  0.0461,\n",
       "           -0.0066, -0.2764,  0.3751,  0.3619,  0.5001, -0.1316],\n",
       "          [-0.6449,  0.5856, -0.0263, -0.0197,  0.8357, -0.5856,  0.0395,\n",
       "           -0.3422,  0.8028,  0.0855, -0.7238, -0.6317,  0.2764, -0.0461,\n",
       "           -0.4211, -0.5988,  0.2632,  0.4014, -0.7501, -0.5659]]],\n",
       "        grad_fn=<UnsqueezeBackward0>), scale=tensor(0.0066, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True)))"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "quant_rnn(torch.randn(2, 5, 10), torch.randn(1, 2, 20))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As with other Brevitas layers, `QuantRNN` can be initialized from a pretrained floating-point `torch.nn.RNN`. For the purpose of this tutorial, can simulate it from an untrained `torch.nn.RNN`. As for other quantized layers, setting `brevitas.config.IGNORE_MISSING_KEYS` might be necessary (depending on which quantizers are set). With the default quantizers, an error on activation scale keys would be triggered, so we set it to true:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:51.544441Z",
     "iopub.status.busy": "2025-05-09T14:50:51.543753Z",
     "iopub.status.idle": "2025-05-09T14:50:51.642495Z",
     "shell.execute_reply": "2025-05-09T14:50:51.640320Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torch.nn import RNN\n",
    "from brevitas.nn import QuantRNN\n",
    "from brevitas import config\n",
    "\n",
    "config.IGNORE_MISSING_KEYS = True\n",
    "\n",
    "float_rnn = RNN(input_size=10, hidden_size=20)\n",
    "quant_rnn = QuantRNN(input_size=10, hidden_size=20)\n",
    "quant_rnn.load_state_dict(float_rnn.state_dict())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similar to other quantized layers, quantization on a certain tensor can be disabled by setting a quantizer to `None`. Setting all quantizers to `None` recovers the same behaviour as the floating-point variant:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:51.651538Z",
     "iopub.status.busy": "2025-05-09T14:50:51.650854Z",
     "iopub.status.idle": "2025-05-09T14:50:51.885259Z",
     "shell.execute_reply": "2025-05-09T14:50:51.882911Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "from torch.nn import RNN\n",
    "from brevitas.nn import QuantRNN\n",
    "from brevitas import config\n",
    "ATOL = 1e-6\n",
    "\n",
    "config.IGNORE_MISSING_KEYS = True\n",
    "torch.manual_seed(123456)\n",
    "\n",
    "float_rnn = RNN(input_size=10, hidden_size=20)\n",
    "quant_rnn = QuantRNN(input_size=10, hidden_size=20, weight_quant=None, io_quant=None, gate_acc_quant=None, bias_quant=None)\n",
    "\n",
    "# Set both layers to the same state_dict\n",
    "quant_rnn.load_state_dict(float_rnn.state_dict())\n",
    "\n",
    "# Generate random input\n",
    "inp = torch.randn(5, 2, 10)\n",
    "# Check outputs are the same\n",
    "assert_with_message(torch.allclose(quant_rnn(inp)[0], float_rnn(inp)[0], atol=ATOL))\n",
    "# Check hidden states are the same\n",
    "assert_with_message(torch.allclose(quant_rnn(inp)[1], float_rnn(inp)[1], atol=ATOL))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As with other quantized layers, we can leverage other prebuilt quantizers too. For example, to perform binary weight quantization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:51.895065Z",
     "iopub.status.busy": "2025-05-09T14:50:51.894240Z",
     "iopub.status.idle": "2025-05-09T14:50:51.988988Z",
     "shell.execute_reply": "2025-05-09T14:50:51.986681Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[-0.3684, -0.0946, -0.4480,  0.0050,  0.1543,  0.6322,  0.1643,\n",
       "            0.1693,  0.2937,  0.5227,  0.2290, -0.3534, -0.3883,  0.4331,\n",
       "            0.0000,  0.1693, -0.4331,  0.3634, -0.0050,  0.1941],\n",
       "          [-0.2240, -0.0199, -0.3534,  0.0946,  0.3485,  0.3534,  0.1941,\n",
       "            0.1643,  0.1145,  0.4082,  0.2987, -0.0647, -0.0946,  0.1543,\n",
       "            0.1145, -0.0498,  0.0647,  0.1493,  0.0299, -0.1195]],\n",
       " \n",
       "         [[ 0.0776, -0.0776, -0.5670,  0.4178, -0.0239,  0.4476,  0.2029,\n",
       "           -0.0836,  0.3521,  0.7042,  0.6326,  0.4058, -0.4118, -0.0477,\n",
       "           -0.2387, -0.0179, -0.4416, -0.4237, -0.3282, -0.1074],\n",
       "          [-0.2626,  0.3581,  0.2328, -0.2268, -0.2686, -0.3103,  0.4536,\n",
       "            0.3461,  0.3103,  0.3163,  0.3282, -0.3163, -0.7639,  0.0179,\n",
       "            0.0060,  0.0776, -0.5849, -0.5252,  0.1790,  0.2984]],\n",
       " \n",
       "         [[-0.5411,  0.3147,  0.6184, -0.3037, -0.1877, -0.3755,  0.1767,\n",
       "           -0.1767, -0.1491, -0.1049,  0.2871, -0.0552, -0.0883,  0.0331,\n",
       "            0.4749, -0.3147,  0.0331,  0.1767,  0.7013, -0.2264],\n",
       "          [-0.0773, -0.1877,  0.4749, -0.2264, -0.4583,  0.0166, -0.3534,\n",
       "           -0.5743,  0.5411,  0.1160, -0.0442, -0.0442,  0.3037,  0.0166,\n",
       "           -0.1325, -0.1657, -0.0718,  0.1215,  0.6240,  0.3092]],\n",
       " \n",
       "         [[-0.0627, -0.1882, -0.4642, -0.1443,  0.4705,  0.3137, -0.2447,\n",
       "            0.0063, -0.1129,  0.3011,  0.1882,  0.2572,  0.2384, -0.0376,\n",
       "            0.1129, -0.1380,  0.1380,  0.3011, -0.0251, -0.0063],\n",
       "          [-0.6399,  0.5771,  0.2133,  0.2572,  0.7967,  0.1631, -0.2384,\n",
       "           -0.4078, -0.3199,  0.0753,  0.6524,  0.0690, -0.1819, -0.2258,\n",
       "            0.3889, -0.4078, -0.3764,  0.2258,  0.5458, -0.1756]],\n",
       " \n",
       "         [[-0.5704,  0.6139, -0.1209, -0.5173,  0.4447,  0.0048,  0.3481,\n",
       "           -0.5946, -0.5221,  0.1644, -0.2949, -0.1789, -0.1982,  0.2707,\n",
       "            0.2900, -0.5124, -0.4399, -0.0725,  0.4351,  0.6091],\n",
       "          [ 0.0435,  0.2030, -0.4447, -0.2659,  0.1547,  0.0580,  0.4254,\n",
       "            0.5559,  0.1740,  0.4254,  0.4592,  0.2369, -0.4496, -0.3336,\n",
       "            0.3046,  0.1354, -0.3626, -0.2659, -0.2079, -0.4641]]],\n",
       "        grad_fn=<CatBackward0>),\n",
       " tensor([[[-0.5704,  0.6139, -0.1209, -0.5173,  0.4447,  0.0048,  0.3481,\n",
       "           -0.5946, -0.5221,  0.1644, -0.2949, -0.1789, -0.1982,  0.2707,\n",
       "            0.2900, -0.5124, -0.4399, -0.0725,  0.4351,  0.6091],\n",
       "          [ 0.0435,  0.2030, -0.4447, -0.2659,  0.1547,  0.0580,  0.4254,\n",
       "            0.5559,  0.1740,  0.4254,  0.4592,  0.2369, -0.4496, -0.3336,\n",
       "            0.3046,  0.1354, -0.3626, -0.2659, -0.2079, -0.4641]]],\n",
       "        grad_fn=<UnsqueezeBackward0>))"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from brevitas.quant.binary import SignedBinaryWeightPerTensorConst\n",
    "\n",
    "binary_rnn = QuantRNN(input_size=10, hidden_size=20, weight_quant=SignedBinaryWeightPerTensorConst)\n",
    "binary_rnn(torch.randn(5, 2, 10))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## QuantLSTM\n",
    "\n",
    "We now look at `QuantLSTM`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:51.998177Z",
     "iopub.status.busy": "2025-05-09T14:50:51.997452Z",
     "iopub.status.idle": "2025-05-09T14:50:52.019598Z",
     "shell.execute_reply": "2025-05-09T14:50:52.017208Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```python\n",
       "    def __init__(\n",
       "            self,\n",
       "            input_size: int,\n",
       "            hidden_size: int,\n",
       "            num_layers: int = 1,\n",
       "            bias: Optional[bool] = True,\n",
       "            batch_first: bool = False,\n",
       "            bidirectional: bool = False,\n",
       "            weight_quant=Int8WeightPerTensorFloat,\n",
       "            bias_quant=Int32Bias,\n",
       "            io_quant=Int8ActPerTensorFloat,\n",
       "            gate_acc_quant=Int8ActPerTensorFloat,\n",
       "            sigmoid_quant=Uint8ActPerTensorFloat,\n",
       "            tanh_quant=Int8ActPerTensorFloat,\n",
       "            cell_state_quant=Int8ActPerTensorFloat,\n",
       "            coupled_input_forget_gates: bool = False,\n",
       "            cat_output_cell_states=True,\n",
       "            shared_input_hidden_weights=False,\n",
       "            shared_intra_layer_weight_quant=False,\n",
       "            shared_intra_layer_gate_acc_quant=False,\n",
       "            shared_cell_state_quant=True,\n",
       "            return_quant_tensor: bool = False,\n",
       "            device: Optional[torch.device] = None,\n",
       "            dtype: Optional[torch.dtype] = None,\n",
       "            **kwargs):\n",
       "        super(QuantLSTM, self).__init__(\n",
       "            layer_impl=_QuantLSTMLayer,\n",
       "            input_size=input_size,\n",
       "            hidden_size=hidden_size,\n",
       "            num_layers=num_layers,\n",
       "            bias=bias,\n",
       "            batch_first=batch_first,\n",
       "            bidirectional=bidirectional,\n",
       "            weight_quant=weight_quant,\n",
       "            bias_quant=bias_quant,\n",
       "            io_quant=io_quant,\n",
       "            gate_acc_quant=gate_acc_quant,\n",
       "            sigmoid_quant=sigmoid_quant,\n",
       "            tanh_quant=tanh_quant,\n",
       "            cell_state_quant=cell_state_quant,\n",
       "            cifg=coupled_input_forget_gates,\n",
       "            shared_input_hidden_weights=shared_input_hidden_weights,\n",
       "            shared_intra_layer_weight_quant=shared_intra_layer_weight_quant,\n",
       "            shared_intra_layer_gate_acc_quant=shared_intra_layer_gate_acc_quant,\n",
       "            shared_cell_state_quant=shared_cell_state_quant,\n",
       "            return_quant_tensor=return_quant_tensor,\n",
       "            dtype=dtype,\n",
       "            device=device,\n",
       "            **kwargs)\n",
       "        if cat_output_cell_states and cell_state_quant is not None and not shared_cell_state_quant:\n",
       "            raise RuntimeError(\"Concatenating cell states requires shared cell quantizers.\")\n",
       "        if return_quant_tensor and cell_state_quant is None:\n",
       "            raise RuntimeError(\"return_quant_tensor=True requires cell_state_quant != None.\")\n",
       "        self.cat_output_cell_states = cat_output_cell_states\n",
       "\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import inspect\n",
    "from brevitas.nn import QuantLSTM\n",
    "from IPython.display import Markdown, display\n",
    "\n",
    "def pretty_print_source(source):\n",
    "    display(Markdown('```python\\n' + source + '\\n```'))\n",
    "    \n",
    "source = inspect.getsource(QuantLSTM.__init__)  \n",
    "pretty_print_source(source)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As with `QuantRNN`, `QuantLSTM` supports all options of `torch.nn.LSTM`. Everything said so far on `QuantRNN` applies to `QuantLSTM` too, but there a bunch of things more to be aware of.\n",
    "\n",
    "`QuantLSTM` accepts a few more quantizers: `sigmoid_quant`, `tanh_quant` and `cell_state_quant`. As with `QuantRNN`, setting `bidirectional=True` and/or `num_layers > 1` triggers sharing the instance of certain quantizers, but not others. In particular `io_quant` is shared among all layers and directions, as it was the case for `QuantRNN`. `cell_state_quant` is shared by default, but setting `shared_cell_state_quant=False` can disable that. However, that requires setting `cat_output_cell_states=False`, as otherwise we would find ourselves with a concenation of cell states that have been quantized with different quantizers, which is considered illegal in Brevitas.\n",
    "\n",
    "LSTMs have four gates, each with its input-hidden and hidden-hidden weights. Brevitas takes in one `weight_quant` definition, but then four different instances of the weight quantizer are instantiated, and each gate is quantized differently, meaning it can have its own scale and zero-point. To force sharing the same weight quantizer across all gates, `QuantLSTM` supports setting `shared_intra_layer_weight_quant=True`. The same reasoning applies to the quantization of the output of each gate, before the activation functions, which is controlled by the `gate_acc_quant` quantizer. To force the same quantizer instance to be shared, `shared_intra_layer_gate_acc_quant=True` can be set. Different sigmoid and tanh functions instead are always allocated different quantizer instances.\n",
    "\n",
    "Finally, `QuantLSTM` also supports the coupled input-forget gates (CIFG), where the forget gate is defined as `forget_gate = 1 - input_gate`, by setting `coupled_input_forget_gates=True`. This is an optimization to save on some compute and number of parameters, and is orthogonal to all other settings, such as `shared_input_hidden_weights`."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Just-in-time compilation\n",
    "\n",
    "Custom recurrent layer can be quite slow at training time. With quantization added in, it only gets worse. To mitigate the issue, both `QuantRNN` and `QuantLSTM` support jit compilation. Setting the env variable `BREVITAS_JIT=1` triggers end-to-end compilation of the quantized recurrent cell through PyTorch TorchScript compiler."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calibration\n",
    "\n",
    "As of version 0.8 of Brevitas, `QuantRNN` and `QuantLSTM` don't support quantized activations calibration through `calibration_mode `nor bias correction through `bias_correction_mode`. This will be added in a future version."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export\n",
    "\n",
    "As of Brevitas 0.8, export of quantized recurrent layers is still a work in progress. As a proof of concept, there is partial support only for export of `QuantLSTM` to ONNX QCDQ, a way to represented quantization in ONNX only with standard ops (QuantizeLinear->Clip->DequantizeLinear), and to QONNX, a custom set of quantized operators introduced by Brevitas on top of ONNX. Two use cases are supported: (1) only `weight_quant` is set, supported by both QCDQ and QONNX, and (2) all quantizers are set, supported only by QONNX. In both cases, `bidirectional=True` and `num_layers > 1` are supported. We first define an utility function to visualize the network through netron, which requires `pip install netron`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:52.030229Z",
     "iopub.status.busy": "2025-05-09T14:50:52.029521Z",
     "iopub.status.idle": "2025-05-09T14:50:52.043012Z",
     "shell.execute_reply": "2025-05-09T14:50:52.040879Z"
    }
   },
   "outputs": [],
   "source": [
    "import time\n",
    "from IPython.display import IFrame\n",
    "\n",
    "def show_netron(model_path, port):\n",
    "    try:\n",
    "        import netron\n",
    "        time.sleep(3.)\n",
    "        netron.start(model_path, address=(\"localhost\", port), browse=False)\n",
    "        return IFrame(src=f\"http://localhost:{port}/\", width=\"100%\", height=400)\n",
    "    except:\n",
    "        pass"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### QuantLSTM weight-only quantization export\n",
    "\n",
    "For use case (1), we leverage export to ONNX QCDQ. Qeight quantization is represented with QCDQ nodes, while the standard ONNX `LSTM` operator is adopted for the recurrent cell. With this approach, we can represent any weight bit width >= 2. Opset 14 is required. For the purpose of this 1 layer, 1 direction example we keep the default `weight_quant` set, we add `weight_bit_width=4`, while we disable the other quantizers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:52.052104Z",
     "iopub.status.busy": "2025-05-09T14:50:52.051401Z",
     "iopub.status.idle": "2025-05-09T14:50:53.074707Z",
     "shell.execute_reply": "2025-05-09T14:50:53.072448Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from brevitas.nn import QuantLSTM\n",
    "from brevitas.export import export_onnx_qcdq\n",
    "\n",
    "quant_lstm_weight_only = QuantLSTM(input_size=10, hidden_size=20, weight_bit_width=4, io_quant=None, bias_quant=None, gate_acc_quant=None, sigmoid_quant=None, tanh_quant=None, cell_state_quant=None)\n",
    "export_path = 'quant_lstm_weight_only_4b.onnx'\n",
    "exported_model = export_onnx_qcdq(quant_lstm_weight_only, (torch.randn(5, 1, 10)), opset_version=14, export_path=export_path)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Serving 'quant_lstm_weight_only_4b.onnx' at http://localhost:8080\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"400\"\n",
       "            src=\"http://localhost:8080/\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x7faa79c37310>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_netron(export_path, 8080)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the model can then be accelerated in `onnxruntime`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:53.088036Z",
     "iopub.status.busy": "2025-05-09T14:50:53.087534Z",
     "iopub.status.idle": "2025-05-09T14:50:53.194475Z",
     "shell.execute_reply": "2025-05-09T14:50:53.192601Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-05-09 15:50:53.160701448 [W:onnxruntime:, graph.cc:1283 Graph] Initializer onnx::LSTM_93 appears in graph inputs and will not be treated as constant value/weight. This may prevent some of the graph optimizations, like const folding. Move it out of graph inputs if there is no need to override it, by either re-generating the model with latest exporter/converter or with the tool onnxruntime/tools/python/remove_initializer_from_input.py.\n"
     ]
    }
   ],
   "source": [
    "import onnxruntime as ort\n",
    "import numpy as np\n",
    "\n",
    "sess = ort.InferenceSession(export_path)\n",
    "input_name = sess.get_inputs()[0].name\n",
    "np_input = np.random.uniform(size=(5, 1, 10)).astype(np.float32)  # (seq_len, batch_size, input_size)\n",
    "pred_onnx = sess.run(None, {input_name: np_input})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "CIFG is also supported in a way that follows the semantics of `onnxruntime`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:53.202410Z",
     "iopub.status.busy": "2025-05-09T14:50:53.202017Z",
     "iopub.status.idle": "2025-05-09T14:50:53.666933Z",
     "shell.execute_reply": "2025-05-09T14:50:53.664732Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from brevitas.nn import QuantLSTM\n",
    "from brevitas.export import export_onnx_qcdq\n",
    "\n",
    "quant_lstm_weight_only_cifg = QuantLSTM(\n",
    "    input_size=10, hidden_size=20, coupled_input_forget_gates=True, weight_bit_width=4,\n",
    "    io_quant=None, bias_quant=None, gate_acc_quant=None, sigmoid_quant=None, tanh_quant=None, cell_state_quant=None)\n",
    "export_path = 'quant_lstm_weight_only_cifg_4b.onnx'\n",
    "exported_model = export_onnx_qcdq(quant_lstm_weight_only_cifg, (torch.randn(5, 1, 10)), opset_version=14, export_path=export_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Serving 'quant_lstm_weight_only_cifg_4b.onnx' at http://localhost:8082\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"400\"\n",
       "            src=\"http://localhost:8082/\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x7faa7ad26650>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_netron(export_path, 8082)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As before we can run it with `onnxruntime`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:53.677210Z",
     "iopub.status.busy": "2025-05-09T14:50:53.676484Z",
     "iopub.status.idle": "2025-05-09T14:50:53.740739Z",
     "shell.execute_reply": "2025-05-09T14:50:53.737604Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-05-09 15:50:53.687943522 [W:onnxruntime:, graph.cc:1283 Graph] Initializer onnx::LSTM_87 appears in graph inputs and will not be treated as constant value/weight. This may prevent some of the graph optimizations, like const folding. Move it out of graph inputs if there is no need to override it, by either re-generating the model with latest exporter/converter or with the tool onnxruntime/tools/python/remove_initializer_from_input.py.\n"
     ]
    }
   ],
   "source": [
    "import onnxruntime as ort\n",
    "import numpy as np\n",
    "\n",
    "sess = ort.InferenceSession(export_path)\n",
    "input_name = sess.get_inputs()[0].name\n",
    "np_input = np.random.uniform(size=(5, 1, 10)).astype(np.float32)  # (seq_len, batch_size, input_size)\n",
    "pred_onnx = sess.run(None, {input_name: np_input})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the 2 layers, 2 directions use case:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:53.749653Z",
     "iopub.status.busy": "2025-05-09T14:50:53.748954Z",
     "iopub.status.idle": "2025-05-09T14:50:55.640637Z",
     "shell.execute_reply": "2025-05-09T14:50:55.638958Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/nn/mixin/base.py:59: UserWarning: Keyword arguments are being passed but they not being used.\n",
      "  warn('Keyword arguments are being passed but they not being used.')\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from brevitas.nn import QuantLSTM\n",
    "from brevitas.export import export_onnx_qcdq\n",
    "\n",
    "quant_lstm_weight_only_bidirectional_2_layers = QuantLSTM(\n",
    "    input_size=10, hidden_size=20, bidirectional=True, num_layers=2, weight_bit_width=4,\n",
    "    io_quant=None, bias_quant=None, gate_acc_quant=None, sigmoid_quant=None, tanh_quant=None, cell_state_quant=None)\n",
    "export_path = 'quant_lstm_weight_only_bidirectional_2_layers.onnx'\n",
    "exported_model = export_onnx_qcdq(quant_lstm_weight_only_bidirectional_2_layers, (torch.randn(5, 1, 10)), opset_version=14, export_path=export_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Serving 'quant_lstm_weight_only_bidirectional_2_layers.onnx' at http://localhost:8083\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"400\"\n",
       "            src=\"http://localhost:8083/\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x7faab66e2150>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_netron(export_path, 8083)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Shared input-hidden weights are also supported:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:55.650365Z",
     "iopub.status.busy": "2025-05-09T14:50:55.649974Z",
     "iopub.status.idle": "2025-05-09T14:50:56.355198Z",
     "shell.execute_reply": "2025-05-09T14:50:56.352971Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/nn/mixin/base.py:59: UserWarning: Keyword arguments are being passed but they not being used.\n",
      "  warn('Keyword arguments are being passed but they not being used.')\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from brevitas.nn import QuantLSTM\n",
    "from brevitas.export import export_onnx_qcdq\n",
    "\n",
    "quant_lstm_weight_only_bidirectional_2_layers_shared = QuantLSTM(\n",
    "    input_size=10, hidden_size=20, bidirectional=True, shared_input_hidden_weights=True, weight_bit_width=4,\n",
    "    io_quant=None, bias_quant=None, gate_acc_quant=None, sigmoid_quant=None, tanh_quant=None, cell_state_quant=None)\n",
    "export_path = 'quant_lstm_weight_only_bidirectional_2_layers_shared_ih.onnx'\n",
    "exported_model = export_onnx_qcdq(quant_lstm_weight_only_bidirectional_2_layers_shared, (torch.randn(5, 1, 10)), opset_version=14, export_path=export_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Serving 'quant_lstm_weight_only_bidirectional_2_layers_shared_ih.onnx' at http://localhost:8085\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"400\"\n",
       "            src=\"http://localhost:8085/\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x7faba0549390>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_netron(export_path, 8085)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can observe how setting `shared_intra_layer_weight_quant=True` affects the network. Now, for each layer and for each direction within a layer, all weight quantizers share the same scale/zp/bit-width:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:56.364956Z",
     "iopub.status.busy": "2025-05-09T14:50:56.364562Z",
     "iopub.status.idle": "2025-05-09T14:50:57.835212Z",
     "shell.execute_reply": "2025-05-09T14:50:57.833482Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/nn/mixin/base.py:59: UserWarning: Keyword arguments are being passed but they not being used.\n",
      "  warn('Keyword arguments are being passed but they not being used.')\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from brevitas.nn import QuantLSTM\n",
    "from brevitas.export import export_onnx_qcdq\n",
    "\n",
    "quant_lstm_weight_only_bidirectional_2_layers = QuantLSTM(\n",
    "    input_size=10, hidden_size=20, bidirectional=True, num_layers=2, weight_bit_width=4, shared_intra_layer_weight_quant=True,\n",
    "    io_quant=None, bias_quant=None, gate_acc_quant=None, sigmoid_quant=None, tanh_quant=None, cell_state_quant=None)\n",
    "export_path = 'quant_lstm_weight_only_bidirectional_2_layers_shared_q.onnx'\n",
    "exported_model = export_onnx_qcdq(quant_lstm_weight_only_bidirectional_2_layers, (torch.randn(5, 1, 10)), opset_version=14, export_path=export_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Serving 'quant_lstm_weight_only_bidirectional_2_layers_shared_q.onnx' at http://localhost:8086\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"400\"\n",
       "            src=\"http://localhost:8086/\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x7faa79c34710>"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_netron(export_path, 8086)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, if we set both `shared_input_hidden_weights=True` and `shared_intra_layer_weight_quant=True`, the side effect is that all quantizers among both directions in a given layer are gonna have the same scale/zp/bit-width. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:57.847312Z",
     "iopub.status.busy": "2025-05-09T14:50:57.846921Z",
     "iopub.status.idle": "2025-05-09T14:50:59.520474Z",
     "shell.execute_reply": "2025-05-09T14:50:59.517966Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/nn/mixin/base.py:59: UserWarning: Keyword arguments are being passed but they not being used.\n",
      "  warn('Keyword arguments are being passed but they not being used.')\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from brevitas.nn import QuantLSTM\n",
    "from brevitas.export import export_onnx_qcdq\n",
    "\n",
    "quant_lstm_weight_only_bidirectional_2_layers = QuantLSTM(\n",
    "    input_size=10, hidden_size=20, bidirectional=True, num_layers=2, weight_bit_width=4, \n",
    "    shared_input_hidden_weights=True, shared_intra_layer_weight_quant=True,\n",
    "    io_quant=None, bias_quant=None, gate_acc_quant=None, sigmoid_quant=None, tanh_quant=None, cell_state_quant=None)\n",
    "export_path = 'quant_lstm_weight_only_bidirectional_2_layers_shared_q_ih.onnx'\n",
    "exported_model = export_onnx_qcdq(quant_lstm_weight_only_bidirectional_2_layers, (torch.randn(5, 1, 10)), opset_version=14, export_path=export_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Serving 'quant_lstm_weight_only_bidirectional_2_layers_shared_q_ih.onnx' at http://localhost:8087\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"400\"\n",
       "            src=\"http://localhost:8087/\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x7faa78c8acd0>"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_netron(export_path, 8087)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### QuantLSTM full quantization export\n",
    "\n",
    "For use case (2) we export to QONNX. Weight quantization is represented with `Quant` nodes, while a custom quantized LSTM operator `QuantLSTMCell` operator is generated for the recurrent cell. Note that currently `QuantLSTMCell` is not yet supported for execution in the `qonnx` library. In a future version of Brevitas, `QuantLSTMCell` will instead be lowered to a series of standard ops + `Quant` nodes. For the purpose example, we keep all quantizers at default:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:50:59.530988Z",
     "iopub.status.busy": "2025-05-09T14:50:59.530525Z",
     "iopub.status.idle": "2025-05-09T14:51:00.369016Z",
     "shell.execute_reply": "2025-05-09T14:51:00.366857Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from brevitas.nn import QuantLSTM\n",
    "from brevitas.export import export_qonnx\n",
    "\n",
    "quant_lstm = QuantLSTM(input_size=10, hidden_size=20)\n",
    "export_path = 'quant_lstm.onnx'\n",
    "exported_model = export_qonnx(quant_lstm, (torch.randn(5, 1, 10)), export_path=export_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Serving 'quant_lstm.onnx' at http://localhost:8088\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"400\"\n",
       "            src=\"http://localhost:8088/\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x7faa78ab76d0>"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_netron(export_path, 8088)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`QuantLSTMCell` takes the following series of inputs. \n",
    "\n",
    "- quant_input, \n",
    "- quant_hidden_state, \n",
    "- quant_cell_state,\n",
    "- quant_weight_ii, \n",
    "- quant_weight_if,\n",
    "- quant_weight_ic,\n",
    "- quant_weight_io,\n",
    "- quant_weight_hi, \n",
    "- quant_weight_hf,\n",
    "- quant_weight_hc,\n",
    "- quant_weight_ho,\n",
    "- quant_bias_input,\n",
    "- quant_bias_forget,\n",
    "- quant_bias_cell,\n",
    "- quant_bias_output,\n",
    "- output_scale, \n",
    "- output_zero_point, \n",
    "- output_bit_width, \n",
    "- cell_state_scale, \n",
    "- cell_state_zero_point, \n",
    "- cell_state_bit_width, \n",
    "- input_acc_scale, \n",
    "- input_acc_zero_point,\n",
    "- input_acc_bit_width, \n",
    "- forget_acc_scale, \n",
    "- forget_acc_zero_point, \n",
    "- forget_acc_bit_width, \n",
    "- cell_acc_scale, \n",
    "- cell_acc_zero_point, \n",
    "- cell_acc_bit_width, \n",
    "- output_acc_scale, \n",
    "- output_acc_zero_point, \n",
    "- output_acc_bit_width, \n",
    "- input_sigmoid_scale, \n",
    "- input_sigmoid_zero_point, \n",
    "- input_sigmoid_bit_width, \n",
    "- forget_sigmoid_scale, \n",
    "- forget_sigmoid_zero_point, \n",
    "- forget_sigmoid_bit_width, \n",
    "- cell_tanh_scale, \n",
    "- cell_tanh_zero_point, \n",
    "- cell_tanh_bit_width, \n",
    "- output_sigmoid_scale, \n",
    "- output_sigmoid_zero_point, \n",
    "- output_sigmoid_bit_width, \n",
    "- hidden_state_tanh_scale, \n",
    "- hidden_state_tanh_zero_point, \n",
    "- hidden_state_tanh_bit_width"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All previous use cases illustrated for the weight-only quantization scenario are also supported."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "c90b4b2913cf255302969887b74954c34035fee81e7901b217735b8a389ceb78"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
